"""
日志报告生成器。

提供各种报告生成器，用于汇总和展示操作执行情况，
如钩子操作摘要、训练统计等。
"""

from typing import Dict, List, Optional, Tuple
from collections import Counter, defaultdict


class SummaryReporter:
    """
    基础摘要报告器抽象类。

    提供生成摘要报告的基础功能，子类可以实现特定类型的报告。
    """

    def __init__(self):
        """初始化摘要报告器"""
        self.reset()

    def reset(self):
        """重置所有统计数据，子类应覆盖此方法"""
        pass

    def _format_tree_node(
        self, text: str, prefix: str = "", is_last: bool = True, is_root: bool = False
    ) -> str:
        """
        格式化树节点文本，添加树形结构标记

        Args:
            text: 节点文本
            prefix: 前缀（用于缩进和连接线）
            is_last: 是否是当前级别的最后一个节点
            is_root: 是否是根节点

        Returns:
            str: 带有树形结构标记的文本
        """
        if is_root:
            return text

        connector = "└── " if is_last else "├── "
        return f"{prefix}{connector}{text}"

    def _format_tree_branch(
        self, lines: List[str], prefix: str = "", is_last: bool = True
    ) -> List[str]:
        """
        格式化树分支，对子节点添加适当的前缀

        Args:
            lines: 子节点文本行列表
            prefix: 当前节点的前缀
            is_last: 当前节点是否是其层级的最后一个节点

        Returns:
            List[str]: 格式化后的树分支文本行
        """
        if not lines:
            return []

        new_prefix = prefix + ("    " if is_last else "│   ")
        return [self._format_tree_node(lines[0], prefix, is_last, False)] + [
            f"{new_prefix}{line}" for line in lines[1:]
        ]

    def _build_tree_report(
        self, nodes: List[Tuple[str, List[str], bool]], title: str
    ) -> List[str]:
        """
        构建树形结构报告

        Args:
            nodes: (节点标题, 子节点文本行, 是否展开子节点) 元组列表
            title: 报告标题

        Returns:
            List[str]: 格式化的树形报告行
        """
        report_lines = [title]

        for i, (node_title, child_lines, expand) in enumerate(nodes):
            is_last = i == len(nodes) - 1
            node_line = self._format_tree_node(node_title, "", is_last, False)
            report_lines.append(node_line)

            if expand and child_lines:
                prefix = "    " if is_last else "│   "
                for j, line in enumerate(child_lines):
                    is_last_child = j == len(child_lines) - 1
                    child_line = self._format_tree_node(
                        line, prefix, is_last_child, False
                    )
                    report_lines.append(child_line)

        return report_lines

    def get_summary_report(self, tree_format: bool = True) -> str:
        """
        生成摘要报告。

        Args:
            tree_format: 是否使用树形格式输出报告

        Returns:
            str: 格式化的摘要报告
        """
        return "摘要报告"


class HookSummaryReporter(SummaryReporter):
    """
    用于生成钩子操作的简洁摘要报告。

    此类收集关于钩子操作（附加、分离）的统计数据，
    并生成结构化、易读的摘要报告。
    """

    def __init__(self):
        """初始化摘要报告器"""
        super().__init__()

    def reset(self):
        """重置所有计数器和统计数据"""
        self.hook_types = Counter()
        self.module_types = Counter()
        self.success_count = 0
        self.failure_count = 0
        self.errors = Counter()
        self.layer_stats = defaultdict(int)
        self.module_paths = defaultdict(list)  # 存储模块路径，用于构建树形结构
        self.recorded_paths = set()  # 用于跟踪已记录的路径，避免重复计数

    def record_hook_creation(
        self,
        hook_type: str,
        module_type: str,
        path: str,
        success: bool,
        error: Optional[str] = None,
    ):
        """
        记录钩子创建的统计数据

        Args:
            hook_type: 钩子类型名称
            module_type: 模块类型名称
            path: 模块路径
            success: 是否成功创建
            error: 错误信息（如果失败）
        """
        if success:
            # 检查该路径是否已经记录过，避免重复计数
            if path in self.recorded_paths:
                return
            self.recorded_paths.add(path)

            self.hook_types[hook_type] += 1
            self.module_types[module_type] += 1
            self.success_count += 1

            # 记录层级统计 (如layer.1, layer.2)
            path_parts = path.split(".")
            for i, part in enumerate(path_parts):
                if (
                    part.startswith("layer")
                    and i + 1 < len(path_parts)
                    and path_parts[i + 1].isdigit()
                ):
                    layer_key = f"{part}.{path_parts[i + 1]}"
                    self.layer_stats[layer_key] += 1
                    break

            # 记录完整的模块路径，用于构建树形结构
            # 将路径按组件分割并存储
            path_components = path.split(".")
            if len(path_components) > 0:
                root = path_components[0]
                self.module_paths[root].append((path, module_type, hook_type))
        else:
            self.failure_count += 1
            if error:
                self.errors[error] += 1

    def _build_model_tree(self) -> Dict[str, List[str]]:
        """
        构建模型结构树

        将记录的模块路径组织成树形结构，按层级关系分组。
        每个节点包含其子节点信息，以及Hook和模块类型信息。

        Returns:
            Dict[str, List[str]]: 模型结构树
        """
        # 按层级组织模型结构
        tree = {}

        # 按路径长度排序，确保父节点先处理
        all_paths = []
        for root_name, paths in self.module_paths.items():
            for full_path, module_type, hook_type in paths:
                all_paths.append((full_path, module_type, hook_type))

        all_paths.sort(key=lambda x: (len(x[0].split(".")), x[0]))

        # 构建树形结构
        for full_path, module_type, hook_type in all_paths:
            components = full_path.split(".")
            current_level = tree

            # 构建完整路径
            for i, component in enumerate(components):
                # 对于路径中的每个组件
                if i < len(components) - 1:  # 非叶子节点
                    if component not in current_level:
                        current_level[component] = {}
                    current_level = current_level[component]
                else:  # 叶子节点（最后一个组件）
                    # 存储带有类型信息的叶子节点
                    current_level[f"{component} [{module_type}] - {hook_type}"] = {}

        return tree

    def _format_model_tree_nodes(self, tree: Dict, prefix: str = "") -> List[str]:
        """
        将树形结构格式化为文本行列表

        Args:
            tree: 树形结构字典
            prefix: 当前节点的前缀

        Returns:
            List[str]: 格式化的文本行列表
        """
        lines = []
        items = list(tree.items())

        # 添加数字值排序，处理层号排序问题
        def sort_key(item):
            key = item[0]
            # 检查键是否只包含数字
            if key.isdigit():
                return int(key)
            # 检查键是否以数字开头但包含其他内容（如 "0 [Linear] - LoRALayerHook"）
            parts = key.split(" ", 1)
            if parts[0].isdigit():
                return int(parts[0])
            return key

        # 对items进行排序，保证层号按数字顺序排序
        try:
            items.sort(key=sort_key)
        except (TypeError, ValueError):
            # 如果排序失败，回退到普通排序
            items.sort()

        for i, (node, subtree) in enumerate(items):
            is_last = i == len(items) - 1
            # 添加当前节点
            connector = "└── " if is_last else "├── "
            lines.append(f"{prefix}{connector}{node}")

            # 添加子节点
            if subtree:
                extension = "    " if is_last else "│   "
                child_lines = self._format_model_tree_nodes(subtree, prefix + extension)
                lines.extend(child_lines)

        return lines

    def get_summary_report(self, tree_format: bool = True) -> str:
        """
        生成简洁的摘要报告

        Args:
            tree_format: 是否使用树形格式输出报告

        Returns:
            str: 格式化的摘要报告
        """
        total = self.success_count + self.failure_count
        success_rate = (self.success_count / total * 100) if total > 0 else 0

        if tree_format:
            # 基本统计节点
            basic_nodes = [
                (
                    "钩子类型分布",
                    [
                        f"{hook_type}: {count}个"
                        for hook_type, count in self.hook_types.most_common()
                    ],
                    True,
                ),
                (
                    "模块类型分布",
                    [
                        f"{module_type}: {count}个"
                        for module_type, count in self.module_types.most_common(5)
                    ],
                    True,
                ),
                (
                    "层级分布",
                    [
                        f"{layer}: {count}个"
                        # 使用自定义排序键函数解决层级排序问题
                        for layer, count in sorted(
                            self.layer_stats.items(),
                            key=lambda item: (
                                int(item[0].split(".")[-1])
                                if item[0].split(".")[-1].isdigit()
                                else item[0]
                            ),
                        )
                    ],
                    True,
                ),
                (
                    "错误摘要",
                    [
                        f"{count}个错误: {error[:77] + '...' if len(error) > 80 else error}"
                        for error, count in self.errors.most_common()
                    ],
                    True,
                ),
            ]

            # 生成基本统计报告
            basic_report = self._build_tree_report(
                basic_nodes,
                f"Hook操作摘要: 成功 {self.success_count}/{total} ({success_rate:.1f}%)",
            )

            # 生成模型结构树报告
            model_tree = self._build_model_tree()
            if model_tree:
                model_tree_lines = ["", "模型结构树 (Hook挂载位置):"]
                formatted_tree = self._format_model_tree_nodes(model_tree)
                model_tree_lines.extend(formatted_tree)

                # 合并两种报告
                return "\n".join(basic_report + model_tree_lines)
            else:
                return "\n".join(basic_report)

        # 非树形格式的输出
        report_lines = [
            f"Hook操作摘要: 成功 {self.success_count}/{total} ({success_rate:.1f}%)"
        ]

        # 添加钩子类型分布
        if self.hook_types:
            report_lines.append("\n钩子类型分布:")
            for hook_type, count in self.hook_types.most_common():
                report_lines.append(f"  - {hook_type}: {count}个")

        # 添加模块类型分布
        if self.module_types:
            report_lines.append("\n模块类型分布:")
            for module_type, count in self.module_types.most_common(5):  # 只显示前5个
                report_lines.append(f"  - {module_type}: {count}个")

        # 如果有层级统计，显示它们
        if self.layer_stats:
            report_lines.append("\n层级分布:")

            # 定义排序键函数，提取层号并转换为整数
            def layer_sort_key(item):
                layer_key = item[0]  # 例如 "layer.10"
                # 尝试提取并转换层号
                try:
                    # 假设格式为 "layer.数字"
                    layer_num = layer_key.split(".")[-1]  # 获取最后一部分，如 "10"
                    return int(layer_num)  # 转换为整数用于排序
                except (ValueError, IndexError):
                    return layer_key  # 如果转换失败，返回原始字符串

            # 使用自定义排序键函数进行排序
            for layer, count in sorted(self.layer_stats.items(), key=layer_sort_key):
                report_lines.append(f"  - {layer}: {count}个")

        # 如果有错误，显示它们
        if self.errors:
            report_lines.append("\n错误摘要:")
            for error, count in self.errors.most_common():
                # 截断过长的错误信息
                if len(error) > 80:
                    error = error[:77] + "..."
                report_lines.append(f"  - {count}个错误: {error}")

        # 添加模型结构树
        if self.module_paths:
            report_lines.append("\n模型结构树 (Hook挂载位置):")
            model_tree = self._build_model_tree()
            formatted_tree = self._format_model_tree_nodes(model_tree)
            for line in formatted_tree:
                report_lines.append(f"  {line}")

        return "\n".join(report_lines)


class TrainingStatsReporter(SummaryReporter):
    """
    用于生成训练过程的摘要统计报告。

    收集和汇总训练过程中的指标，如损失、准确率等，并生成结构化报告。
    """

    def __init__(self):
        """初始化训练统计报告器"""
        super().__init__()

    def reset(self):
        """重置统计数据"""
        self.metrics = defaultdict(list)
        self.epoch_metrics = defaultdict(dict)
        self.current_epoch = 0

    def record_metric(self, name: str, value: float, epoch: Optional[int] = None):
        """
        记录单个指标值

        Args:
            name: 指标名称
            value: 指标值
            epoch: 指标所属的轮次，如果为None则使用当前轮次
        """
        epoch = epoch if epoch is not None else self.current_epoch

        # 记录到整体指标
        self.metrics[name].append(value)

        # 记录到特定轮次
        self.epoch_metrics[epoch][name] = value

    def set_current_epoch(self, epoch: int):
        """
        设置当前轮次

        Args:
            epoch: 当前训练轮次
        """
        self.current_epoch = epoch

    def get_summary_report(
        self, include_history: bool = False, tree_format: bool = True
    ) -> str:
        """
        生成训练统计摘要报告

        Args:
            include_history: 是否包含历史记录
            tree_format: 是否使用树形格式输出报告

        Returns:
            str: 格式化的摘要报告
        """
        title = "训练统计摘要"

        if tree_format:
            nodes = []

            # 添加最新轮次的指标
            if self.epoch_metrics:
                latest_epoch = max(self.epoch_metrics.keys())
                latest_metrics = []

                for name, value in sorted(self.epoch_metrics[latest_epoch].items()):
                    latest_metrics.append(f"{name}: {value:.4f}")

                nodes.append((f"轮次 {latest_epoch} 指标", latest_metrics, True))

            # 如果需要，添加历史记录
            if include_history and self.metrics:
                for name, values in sorted(self.metrics.items()):
                    if values:
                        avg = sum(values) / len(values)
                        min_val = min(values)
                        max_val = max(values)

                        history_details = [
                            f"平均值: {avg:.4f}",
                            f"最小值: {min_val:.4f}",
                            f"最大值: {max_val:.4f}",
                            f"样本数: {len(values)}",
                        ]

                        nodes.append((f"指标历史: {name}", history_details, True))

            # 添加当前轮次
            nodes.append((f"当前轮次: {self.current_epoch}", [], False))

            # 按轮次添加所有指标
            epochs = sorted(self.epoch_metrics.keys())
            if epochs:
                epoch_nodes = []
                for epoch in epochs:
                    metric_lines = [
                        f"{name}: {value:.4f}"
                        for name, value in sorted(self.epoch_metrics[epoch].items())
                    ]
                    epoch_nodes.append(
                        (f"轮次 {epoch}", metric_lines, epoch == latest_epoch)
                    )

                nodes.append(("所有轮次数据", [], False))
                nodes.extend(epoch_nodes)

            return "\n".join(self._build_tree_report(nodes, title))
        else:
            report_lines = [title]

            # 添加最新轮次的指标
            if self.epoch_metrics:
                latest_epoch = max(self.epoch_metrics.keys())
                report_lines.append(f"\n轮次 {latest_epoch} 指标:")

                for name, value in sorted(self.epoch_metrics[latest_epoch].items()):
                    report_lines.append(f"  - {name}: {value:.4f}")

            # 如果需要，添加历史记录
            if include_history and self.metrics:
                report_lines.append("\n指标历史统计:")

                for name, values in sorted(self.metrics.items()):
                    if values:
                        avg = sum(values) / len(values)
                        min_val = min(values)
                        max_val = max(values)
                        report_lines.append(
                            f"  - {name}: 平均={avg:.4f}, 最小={min_val:.4f}, 最大={max_val:.4f}"
                        )

            return "\n".join(report_lines)

    def get_current_metrics(self) -> Dict[str, float]:
        """
        获取当前轮次的所有指标

        Returns:
            Dict[str, float]: 指标名称到值的映射
        """
        return self.epoch_metrics.get(self.current_epoch, {})
